// Based on example from: http://hax.fi/asko/PSSM.html, Thanks for Jeroen for him hard work.

#include "PSSMUtils.h"

IrrlichtDevice* Device = 0;
IVideoDriver* Driver = 0;
ISceneManager* SceneManager = 0;
ICameraSceneNode* Camera = 0;
ICameraSceneNode* LightCamera = 0;
stringw window = "";

float CameraNear = 1.0f;
float CameraFar = 500.0f;

float LightCameraNear = 150.0f;
float LightCameraFar = 400.0f;

class CDepthCallBack: public ICgShaderConstantSetCallBack
{
public:
    CDepthCallBack(CPSSMUtils* vPSSMUtils) : PSSMUtils(vPSSMUtils), WVP(0), WV(0)
    {
    }

    ~CDepthCallBack(){};

	virtual void OnSetConstants(ICgServices* Services,const CGeffect& Effect,const CGpass& Pass,const CGprogram& Vertex,const CGprogram& Fragment,const SMaterial& Material)
	{
	    matrix4 Matrix;

        WVP = cgGetNamedParameter(Vertex, "WVP");
		Matrix = PSSMUtils->getProjectionMatrix(PSSMUtils->getCurrentPass());
		Matrix *= PSSMUtils->getViewMatrix(PSSMUtils->getCurrentPass());
		Matrix *= Driver->getTransform(ETS_WORLD);
		Services->setMatrix( WVP,ICGT_MATRIX_IDENTITY,Matrix );

		WV = cgGetNamedParameter(Vertex, "WV");
		Matrix = PSSMUtils->getViewMatrix(PSSMUtils->getCurrentPass());
		Matrix *= Driver->getTransform(ETS_WORLD);
		Services->setMatrix( WV,ICGT_MATRIX_IDENTITY,Matrix );
	}

private:
    CPSSMUtils* PSSMUtils;

    CGparameter	WVP;
    CGparameter	WV;
};

class CPSSMCallBack : public ICgShaderConstantSetCallBack
{
public:
    CPSSMCallBack(CPSSMUtils* vPSSMUtils) : PSSMUtils(vPSSMUtils), vNumSplits(0), WVP(0), TextureMatrix(0), pNumSplits(0),
                                            SplitDistance(0), Bias(0), TextureSize(0), DepthMap(0), CurrentParam(0)
    {
        vBias[0] = 0.001;
        vBias[1] = 0.0008;
        vBias[2] = 0.00049;
        vBias[3] = 0.00015;
    }

    ~CPSSMCallBack(){};

    virtual void OnSetConstants(ICgServices* Services,const CGeffect& Effect,const CGpass& Pass,const CGprogram& Vertex,const CGprogram& Fragment,const SMaterial& Material)
    {
        matrix4 Matrix;

        vNumSplits = cgGetNamedParameter(Vertex, "NumSplits");
        Services->setParameter1i(vNumSplits, PSSMUtils->getSplitsCount());

        WVP = cgGetNamedParameter(Vertex, "WVP");
		Matrix = Driver->getTransform(ETS_PROJECTION);
        Matrix *= Driver->getTransform(ETS_VIEW);
		Matrix *= Driver->getTransform(ETS_WORLD);
		Services->setMatrix(WVP, ICGT_MATRIX_IDENTITY, Matrix);
        WVP = cgGetNamedParameter(Vertex, "WVP");

        pNumSplits = cgGetNamedParameter(Fragment, "NumSplits");
        Services->setParameter1i(pNumSplits, PSSMUtils->getSplitsCount());

        float TextureOffset = 0.5f;

        matrix4 mTextureMatrixConst;
        mTextureMatrixConst[0] = 0.5f;
        mTextureMatrixConst[1] = 0.0f;
        mTextureMatrixConst[2] = 0.0f;
        mTextureMatrixConst[3] = 0.0f;

        mTextureMatrixConst[4] = 0.0f;
        mTextureMatrixConst[5] = -0.5f;

        if(Driver->getDriverType() == EDT_OPENGL)
        mTextureMatrixConst[5] *= -1.0f;

        mTextureMatrixConst[6] = 0.0f;
        mTextureMatrixConst[7] = 0.0f;

        mTextureMatrixConst[8] = 0.0f;
        mTextureMatrixConst[9] = 0.0f;
        mTextureMatrixConst[10] = 1.0f;
        mTextureMatrixConst[11] = 0.0f;

        mTextureMatrixConst[12] = TextureOffset;
        mTextureMatrixConst[13] = TextureOffset;
        mTextureMatrixConst[14] = 0.0f;
        mTextureMatrixConst[15] = 1.0f;

        TextureMatrix = cgGetNamedParameter(Vertex, "TextureMatrix");
        for(int i = 0; i < PSSMUtils->getSplitsCount(); i++)
        {
            TextureOffset = 0.5f;

            if(Driver->getDriverType() == EDT_OPENGL)
            TextureOffset -= (0.5f / (float)Material.TextureLayer[i].Texture->getSize().Width);
            else
            TextureOffset += (0.5f / (float)Material.TextureLayer[i].Texture->getSize().Width);

            mTextureMatrixConst[12] = TextureOffset;
            mTextureMatrixConst[13] = TextureOffset;

            CurrentParam = cgGetArrayParameter(TextureMatrix, i);

            Matrix = mTextureMatrixConst;
            Matrix *= PSSMUtils->getProjectionMatrix(i);
            Matrix *= PSSMUtils->getViewMatrix(i);
            Matrix *= Driver->getTransform(ETS_WORLD);

            Services->setMatrix( CurrentParam,ICGT_MATRIX_IDENTITY,Matrix );
        }

        SplitDistance = cgGetNamedParameter(Fragment, "SplitDistance");
        for(int i = 0; i < PSSMUtils->getSplitsCount(); i++)
        {
            CurrentParam = cgGetArrayParameter(SplitDistance, i);
            Services->setParameter1f(CurrentParam, PSSMUtils->getSplitDistance(i+1));
        }

        Bias = cgGetNamedParameter(Fragment, "Bias");
        for(int i = 0; i < PSSMUtils->getSplitsCount(); i++)
        {
            CurrentParam = cgGetArrayParameter(Bias, i);
            Services->setParameter1f(CurrentParam, vBias[i]);
        }

        TextureSize = cgGetNamedParameter(Fragment, "TextureSize");
        for(int i = 0; i < PSSMUtils->getSplitsCount(); i++)
        {
            float vTextureSize = (float)Material.TextureLayer[i].Texture->getSize().Width;
            vTextureSize *= (i+1);

            CurrentParam = cgGetArrayParameter(TextureSize, i);
            Services->setParameter1f(CurrentParam, vTextureSize);
        }

        DepthMap = cgGetNamedParameter(Fragment, "DepthMap");
        for(int i = 0; i < PSSMUtils->getSplitsCount(); i++)
        {
            CurrentParam = cgGetArrayParameter(DepthMap, i);
            Services->EnableTexture( CurrentParam,Material.TextureLayer[i].Texture);
        }
	}

private:
    CPSSMUtils* PSSMUtils;
    float vBias[4];

    CGparameter vNumSplits;
    CGparameter	WVP;
    CGparameter TextureMatrix;

    CGparameter pNumSplits;
    CGparameter SplitDistance;
    CGparameter Bias;
    CGparameter TextureSize;
    CGparameter	DepthMap;

    CGparameter CurrentParam;
};


int main()
{
    printf("\n *** PSSM Example by Patryk Nadrowski ***\n\n");

    char i;

    // Select Irrlicht Driver
    video::E_DRIVER_TYPE DriverType;
	printf("\n Please select the driver you want for this demo:\n\n"\
		"  (a) Direct3D 9.0c\n  (b) OpenGL 1.5\n"\
		"  (otherKey) exit\n\n");

	std::cin >> i;

	switch(i)
	{
		case 'a':
            DriverType = EDT_DIRECT3D9;
            window = " Direct3D9 PS: ";
            break;

		case 'b':
            DriverType = EDT_OPENGL;
            window = " OpenGL PS: ";
            break;

		default: return 1;
	}

	int TextureSize = 512;

	printf("\n Please select a textures resolution:\n\n"\
		"  (a) 2048x2048\n  (b) 1024x1024\n  (c) 512x512 (Default)\n");

	std::cin >> i;

	switch(i)
	{
		case 'a':
            TextureSize = 2048;
            break;

		case 'b':
            TextureSize = 1024;
            break;

		default:
            TextureSize = 512;
            break;
	}

	Device = createDevice(DriverType,dimension2d<u32>(800, 600), 32, false, false, false, 0);
	Driver = Device->getVideoDriver();
	SceneManager = Device->getSceneManager();

    // Create Cg Programming Service
    ICgProgrammingServices* GPU = new ICgProgrammingServices(Device);

    // Create PSSM
    LightCamera = SceneManager->addCameraSceneNode();
    LightCamera->setNearValue(LightCameraNear);
    LightCamera->setFarValue(LightCameraFar);
    LightCamera->setFOV(90.0f*(PI/180.0f));
    LightCamera->setPosition(vector3df(-220,140,-212));
    LightCamera->setTarget(vector3df(0,0,0));

    CPSSMUtils* PSSMUtils = new CPSSMUtils(LightCamera, 3, 0.5);

    // Create Materials
    CDepthCallBack* DCB = new CDepthCallBack(PSSMUtils);
    s32 MT_Depth = GPU->addCgShaderMaterialFromFiles(CG_SOURCE,"../../cg/pssm/depth.cg","main_v","arbvp1","vs_2_0","../../cg/pssm/depth.cg","main_f","arbfp1","ps_2_0",DCB,EMT_SOLID);

    CPSSMCallBack* PSSMCB = new CPSSMCallBack(PSSMUtils);
    s32 MT_PSSM = GPU->addCgShaderMaterialFromFiles(CG_SOURCE,"../../cg/pssm/pssm.cg","main_v","arbvp1","hlslv","../../cg/pssm/pssm.cg","main_f","arbfp1","hlslf",PSSMCB,EMT_SOLID);

    // Create Node
    IMesh* mesh = SceneManager->getMesh("../../media/city.x")->getMesh(0);

    mesh->getMeshBuffer(0)->setHardwareMappingHint(EHM_STATIC);

	IMeshSceneNode* node = SceneManager->addMeshSceneNode(mesh);
	node->setPosition(core::vector3df(0,0,0));
	node->setRotation(core::vector3df(0,0,0));
	node->setScale(core::vector3df(1,1,1));
	node->setAutomaticCulling(EAC_FRUSTUM_BOX);
	node->setMaterialFlag(video::EMF_LIGHTING, false);

	// Create Textures
	ITexture* RT0 = Driver->addRenderTargetTexture(dimension2d<u32>(TextureSize, TextureSize),"",ECF_R16F); // First texture is 2x higher
	ITexture* RT1 = Driver->addRenderTargetTexture(dimension2d<u32>(TextureSize, TextureSize),"",ECF_R16F);
	ITexture* RT2 = Driver->addRenderTargetTexture(dimension2d<u32>(TextureSize, TextureSize),"",ECF_R16F);
	ITexture* RT3 = Driver->addRenderTargetTexture(dimension2d<u32>(TextureSize, TextureSize),"",ECF_R16F);

	array<ITexture*> RT;
	RT.push_back(RT0);
	RT.push_back(RT1);
	RT.push_back(RT2);
	RT.push_back(RT3);

	// Apply textures to node
	for(int i = 0; i < 4; i++)
    node->getMaterial(0).TextureLayer[i].Texture = RT[i];

    // Create Camera
    Camera = SceneManager->addCameraSceneNodeFPS(0,100.0f,0.05);
	Camera->setPosition(vector3df(0,20,-110));
	Camera->setTarget(vector3df (0,20,-109));
	Camera->setNearValue(CameraNear);
    Camera->setFarValue(CameraFar);
    Camera->setFOV(45.0f*(PI/180.0f));
    Camera->setUpVector(vector3df (0,1,0));

    node->getMaterial(0).TextureLayer[0].TextureWrapU = ETC_CLAMP_TO_BORDER;
    node->getMaterial(0).TextureLayer[1].TextureWrapU = ETC_CLAMP_TO_BORDER;
    node->getMaterial(0).TextureLayer[2].TextureWrapU = ETC_CLAMP_TO_BORDER;
    node->getMaterial(0).TextureLayer[3].TextureWrapU = ETC_CLAMP_TO_BORDER;

    node->getMaterial(0).TextureLayer[0].TextureWrapV = ETC_CLAMP_TO_BORDER;
    node->getMaterial(0).TextureLayer[1].TextureWrapV = ETC_CLAMP_TO_BORDER;
    node->getMaterial(0).TextureLayer[2].TextureWrapV = ETC_CLAMP_TO_BORDER;
    node->getMaterial(0).TextureLayer[3].TextureWrapV = ETC_CLAMP_TO_BORDER;

    node->getMaterial(0).setFlag(EMF_BILINEAR_FILTER, 1);

    // Last FPS
    int lastFPS = -1;

	while(Device->run())
    if (Device->isWindowActive())
	{
	    // Begin Scene
		Driver->beginScene(true, true, video::SColor(255,0,0,0));

        SceneManager->setActiveCamera(LightCamera);
        PSSMUtils->CalculateSplitDistances(Camera);

        // Render only back faces
        node->setMaterialFlag(EMF_BACK_FACE_CULLING, false);
        node->setMaterialFlag(EMF_FRONT_FACE_CULLING, true);

        for(int CurrentSplit = 0; CurrentSplit < PSSMUtils->getSplitsCount(); CurrentSplit++)
        {
            Driver->setRenderTarget(RT[CurrentSplit], true, true, SColor(255,255,255,255));

            PSSMUtils->CalculateFrustumCorners(Camera,CurrentSplit,1.1);
            PSSMUtils->CalculateLightForFrustum(CurrentSplit);

            node->setMaterialType((E_MATERIAL_TYPE)MT_Depth);

            SceneManager->drawAll();

            Driver->setRenderTarget(0);
        }

        SceneManager->setActiveCamera(Camera);

        node->setMaterialType((E_MATERIAL_TYPE)MT_PSSM);

        node->setMaterialFlag(EMF_BILINEAR_FILTER, false);
        // Render only front faces - standard rendering
        node->setMaterialFlag(video::EMF_BACK_FACE_CULLING, true);
		node->setMaterialFlag(video::EMF_FRONT_FACE_CULLING, false);

        SceneManager->drawAll();

        // End Scene
		Driver->endScene();

		int fps = Driver->getFPS();
		int poly = Driver->getPrimitiveCountDrawn();

		if(lastFPS == -1)
        Camera->setTarget(vector3df (6,-4,-13));

		if (lastFPS != fps)
		{
		    core::stringw str = "";
            str = L"Irrlicht Cg - PSSM for";
            str += window;
            str += L" FPS: ";
            str += fps;
            str += L" Poly: ";
            str += poly;

            Device->setWindowCaption(str.c_str());
            lastFPS = fps;
		}
	}

    delete DCB;
	delete PSSMCB;
	delete GPU;
	delete PSSMUtils;

	Device->drop();

	return 0;
}

